package com.icehand.demo.utils;

import com.google.common.util.concurrent.ThreadFactoryBuilder;

import java.util.concurrent.*;

/**
 * @author zhoujunmei
 */
public class CountDownLatchUtil {
    private CountDownLatch start;
    private CountDownLatch end;
    private int pollSize;

    public CountDownLatchUtil(){
        this(10);
    }

    public CountDownLatchUtil(int pollSize){
        this.pollSize = pollSize;
        start = new CountDownLatch(1);
        end = new CountDownLatch(pollSize);
    }

    @FunctionalInterface
    public interface FunctionInterface{
        void exec();
    }

    public void latch(FunctionInterface functionInterface) throws InterruptedException {
        ThreadFactory threadFactory = new ThreadFactoryBuilder().setNameFormat("latch-util-pool-%s").build();
        ExecutorService executorService = new ThreadPoolExecutor(pollSize, pollSize,
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>(),threadFactory,new ThreadPoolExecutor.AbortPolicy());
        for(int i=0;i<pollSize;i++){
            Runnable run = () -> {
                try {
                    start.await();
                    functionInterface.exec();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }finally {
                    end.countDown();
                }
            };
            executorService.submit(run);
        }

        start.countDown();
        end.await();
        executorService.shutdown();
    }
}
